package com.flink.hbase.statemanagement.checkpoint;

import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.util.Collector;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * CheckpointedFunction接口实现示例
 * 
 * 展示如何实现CheckpointedFunction接口：
 * 1. 状态初始化和恢复
 * 2. 状态快照保存
 * 3. 算子状态管理
 * 4. 自定义状态恢复逻辑
 * 5. 状态一致性保证
 * 
 * @author Flink Team
 */
public class CheckpointedFunctionExample {
    
    /**
     * 计数器事件
     */
    public static class CounterEvent {
        public String key;
        public long value;
        public String type;
        public long timestamp;
        
        public CounterEvent() {}
        
        public CounterEvent(String key, long value, String type, long timestamp) {
            this.key = key;
            this.value = value;
            this.type = type;
            this.timestamp = timestamp;
        }
        
        @Override
        public String toString() {
            return String.format("CounterEvent{key='%s', value=%d, type='%s', timestamp=%d}", 
                key, value, type, timestamp);
        }
    }
    
    /**
     * 累加器状态
     */
    public static class AccumulatorState {
        public String key;
        public long count;
        public long sum;
        public long min;
        public long max;
        public double average;
        public long lastUpdateTime;
        
        public AccumulatorState() {}
        
        public AccumulatorState(String key) {
            this.key = key;
            this.count = 0;
            this.sum = 0;
            this.min = Long.MAX_VALUE;
            this.max = Long.MIN_VALUE;
            this.average = 0.0;
            this.lastUpdateTime = System.currentTimeMillis();
        }
        
        public void update(long value) {
            this.count++;
            this.sum += value;
            this.min = Math.min(this.min, value);
            this.max = Math.max(this.max, value);
            this.average = (double) this.sum / this.count;
            this.lastUpdateTime = System.currentTimeMillis();
        }
        
        @Override
        public String toString() {
            return String.format("AccumulatorState{key='%s', count=%d, sum=%d, min=%d, max=%d, average=%.2f, lastUpdateTime=%d}", 
                key, count, sum, min, max, average, lastUpdateTime);
        }
    }
    
    /**
     * 状态管理器
     * 实现CheckpointedFunction接口，管理算子状态
     */
    public static class StateManager extends RichMapFunction<CounterEvent, Tuple3<String, String, Long>> 
            implements CheckpointedFunction {
        
        // 运行时状态（内存中的状态）
        private Map<String, AccumulatorState> accumulatorStates;
        
        // 用于Checkpoint的算子状态
        private ListState<AccumulatorState> checkpointedState;
        
        // 处理计数器（用于演示状态恢复）
        private long processedCount = 0;
        
        @Override
        public void open(Configuration parameters) throws Exception {
            // 初始化运行时状态
            accumulatorStates = new ConcurrentHashMap<>();
            
            System.out.println("🔧 StateManager opened, 子任务索引: " + 
                getRuntimeContext().getIndexOfThisSubtask());
        }
        
        @Override
        public Tuple3<String, String, Long> map(CounterEvent event) throws Exception {
            // 获取或创建累加器状态
            AccumulatorState state = accumulatorStates.computeIfAbsent(event.key, AccumulatorState::new);
            
            // 更新状态
            state.update(event.value);
            
            // 增加处理计数
            processedCount++;
            
            // 根据事件类型返回不同的结果
            String resultType;
            long resultValue;
            
            switch (event.type) {
                case "COUNT":
                    resultType = "COUNT";
                    resultValue = state.count;
                    break;
                case "SUM":
                    resultType = "SUM";
                    resultValue = state.sum;
                    break;
                case "MIN":
                    resultType = "MIN";
                    resultValue = state.min;
                    break;
                case "MAX":
                    resultType = "MAX";
                    resultValue = state.max;
                    break;
                case "AVG":
                    resultType = "AVG";
                    resultValue = (long) state.average;
                    break;
                default:
                    resultType = "UNKNOWN";
                    resultValue = 0;
            }
            
            System.out.println(String.format("🔄 状态更新: %s -> %s = %d (处理总数: %d)", 
                event.key, resultType, resultValue, processedCount));
            
            return new Tuple3<>(event.key, resultType, resultValue);
        }
        
        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            // 清空之前的状态
            checkpointedState.clear();
            
            // 保存所有累加器状态
            for (AccumulatorState state : accumulatorStates.values()) {
                checkpointedState.add(state);
            }
            
            System.out.println(String.format("📸 状态快照: Checkpoint %d, 保存 %d 个状态, 处理计数: %d", 
                context.getCheckpointId(), accumulatorStates.size(), processedCount));
        }
        
        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            // 创建算子状态描述符
            ListStateDescriptor<AccumulatorState> descriptor = new ListStateDescriptor<>(
                "accumulator-states",
                TypeInformation.of(AccumulatorState.class)
            );
            
            // 获取或创建算子状态
            checkpointedState = context.getOperatorStateStore().getListState(descriptor);
            
            // 如果是从checkpoint恢复，则恢复状态
            if (context.isRestored()) {
                accumulatorStates = new ConcurrentHashMap<>();
                
                // 恢复所有累加器状态
                for (AccumulatorState state : checkpointedState.get()) {
                    accumulatorStates.put(state.key, state);
                }
                
                System.out.println(String.format("🔄 状态恢复: 恢复 %d 个状态", accumulatorStates.size()));
                
                // 输出恢复的状态详情
                for (AccumulatorState state : accumulatorStates.values()) {
                    System.out.println("  恢复状态: " + state);
                }
            } else {
                accumulatorStates = new ConcurrentHashMap<>();
                System.out.println("🆕 初始化新状态");
            }
        }
    }
    
    /**
     * 批处理状态管理器
     * 实现批量状态处理和优化
     */
    public static class BatchStateProcessor extends RichMapFunction<CounterEvent, Tuple2<String, Map<String, Long>>> 
            implements CheckpointedFunction {
        
        // 批处理缓存
        private Map<String, List<CounterEvent>> batchCache;
        
        // 批处理状态
        private ListState<Map<String, List<CounterEvent>>> batchCheckpointedState;
        
        // 批处理大小
        private static final int BATCH_SIZE = 100;
        
        @Override
        public void open(Configuration parameters) throws Exception {
            batchCache = new ConcurrentHashMap<>();
        }
        
        @Override
        public Tuple2<String, Map<String, Long>> map(CounterEvent event) throws Exception {
            // 将事件添加到批处理缓存
            batchCache.computeIfAbsent(event.key, k -> new ArrayList<>()).add(event);
            
            // 检查是否达到批处理大小
            List<CounterEvent> events = batchCache.get(event.key);
            if (events.size() >= BATCH_SIZE) {
                // 处理批次
                Map<String, Long> batchResult = processBatch(event.key, events);
                
                // 清空批处理缓存
                batchCache.remove(event.key);
                
                return new Tuple2<>(event.key, batchResult);
            }
            
            // 如果未达到批处理大小，返回空结果
            return new Tuple2<>(event.key, Collections.emptyMap());
        }
        
        /**
         * 处理批次数据
         */
        private Map<String, Long> processBatch(String key, List<CounterEvent> events) {
            Map<String, Long> result = new HashMap<>();
            
            long count = events.size();
            long sum = events.stream().mapToLong(e -> e.value).sum();
            long min = events.stream().mapToLong(e -> e.value).min().orElse(0);
            long max = events.stream().mapToLong(e -> e.value).max().orElse(0);
            long avg = sum / count;
            
            result.put("count", count);
            result.put("sum", sum);
            result.put("min", min);
            result.put("max", max);
            result.put("avg", avg);
            
            System.out.println(String.format("🔄 批处理完成: %s 处理 %d 个事件", key, count));
            
            return result;
        }
        
        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            // 清空之前的状态
            batchCheckpointedState.clear();
            
            // 保存批处理缓存
            if (!batchCache.isEmpty()) {
                batchCheckpointedState.add(new HashMap<>(batchCache));
            }
            
            System.out.println(String.format("📸 批处理状态快照: Checkpoint %d, 缓存键数: %d", 
                context.getCheckpointId(), batchCache.size()));
        }
        
        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            // 创建算子状态描述符
            ListStateDescriptor<Map<String, List<CounterEvent>>> descriptor = new ListStateDescriptor<>(
                "batch-cache-states",
                TypeInformation.of(Map.class)
            );
            
            // 获取或创建算子状态
            batchCheckpointedState = context.getOperatorStateStore().getListState(descriptor);
            
            // 如果是从checkpoint恢复，则恢复状态
            if (context.isRestored()) {
                batchCache = new ConcurrentHashMap<>();
                
                // 恢复批处理缓存
                for (Map<String, List<CounterEvent>> cache : batchCheckpointedState.get()) {
                    batchCache.putAll(cache);
                }
                
                System.out.println(String.format("🔄 批处理状态恢复: 恢复 %d 个缓存键", batchCache.size()));
            } else {
                batchCache = new ConcurrentHashMap<>();
                System.out.println("🆕 初始化新批处理状态");
            }
        }
    }
    
    /**
     * 状态监控Sink
     * 实现状态监控和统计
     */
    public static class StateMonitorSink extends RichSinkFunction<Tuple3<String, String, Long>> 
            implements CheckpointedFunction {
        
        // 监控统计
        private Map<String, Long> monitorStats;
        
        // 监控状态
        private ListState<Map<String, Long>> monitorCheckpointedState;
        
        @Override
        public void open(Configuration parameters) throws Exception {
            monitorStats = new ConcurrentHashMap<>();
        }
        
        @Override
        public void invoke(Tuple3<String, String, Long> value, Context context) throws Exception {
            // 更新监控统计
            String statKey = value.f0 + "_" + value.f1;
            monitorStats.put(statKey, value.f2);
            
            // 定期输出监控报告
            if (monitorStats.size() % 10 == 0) {
                printMonitorReport();
            }
        }
        
        /**
         * 输出监控报告
         */
        private void printMonitorReport() {
            System.out.println("\n=== 状态监控报告 ===");
            System.out.println("监控时间: " + new Date());
            System.out.println("监控指标数: " + monitorStats.size());
            
            // 按键分组统计
            Map<String, Map<String, Long>> groupedStats = new HashMap<>();
            for (Map.Entry<String, Long> entry : monitorStats.entrySet()) {
                String[] parts = entry.getKey().split("_");
                if (parts.length >= 2) {
                    String key = parts[0];
                    String type = parts[1];
                    groupedStats.computeIfAbsent(key, k -> new HashMap<>()).put(type, entry.getValue());
                }
            }
            
            // 输出分组统计
            for (Map.Entry<String, Map<String, Long>> entry : groupedStats.entrySet()) {
                System.out.println("键: " + entry.getKey() + " -> " + entry.getValue());
            }
            
            System.out.println("==================\n");
        }
        
        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            // 清空之前的状态
            monitorCheckpointedState.clear();
            
            // 保存监控统计
            if (!monitorStats.isEmpty()) {
                monitorCheckpointedState.add(new HashMap<>(monitorStats));
            }
            
            System.out.println(String.format("📸 监控状态快照: Checkpoint %d, 监控指标数: %d", 
                context.getCheckpointId(), monitorStats.size()));
        }
        
        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            // 创建算子状态描述符
            ListStateDescriptor<Map<String, Long>> descriptor = new ListStateDescriptor<>(
                "monitor-stats-states",
                TypeInformation.of(Map.class)
            );
            
            // 获取或创建算子状态
            monitorCheckpointedState = context.getOperatorStateStore().getListState(descriptor);
            
            // 如果是从checkpoint恢复，则恢复状态
            if (context.isRestored()) {
                monitorStats = new ConcurrentHashMap<>();
                
                // 恢复监控统计
                for (Map<String, Long> stats : monitorCheckpointedState.get()) {
                    monitorStats.putAll(stats);
                }
                
                System.out.println(String.format("🔄 监控状态恢复: 恢复 %d 个监控指标", monitorStats.size()));
            } else {
                monitorStats = new ConcurrentHashMap<>();
                System.out.println("🆕 初始化新监控状态");
            }
        }
    }
    
    /**
     * 模拟计数器事件源
     */
    public static class MockCounterEventSource implements SourceFunction<CounterEvent> {
        private volatile boolean running = true;
        private final Random random = new Random();
        private final String[] keys = {"key1", "key2", "key3", "key4", "key5"};
        private final String[] types = {"COUNT", "SUM", "MIN", "MAX", "AVG"};
        
        @Override
        public void run(SourceContext<CounterEvent> ctx) throws Exception {
            while (running) {
                CounterEvent event = new CounterEvent(
                    keys[random.nextInt(keys.length)],
                    random.nextLong() % 100 + 1, // 1-100随机值
                    types[random.nextInt(types.length)],
                    System.currentTimeMillis()
                );
                
                ctx.collect(event);
                Thread.sleep(500 + random.nextInt(1000)); // 0.5-1.5秒随机间隔
            }
        }
        
        @Override
        public void cancel() {
            running = false;
        }
    }
    
    /**
     * 运行CheckpointedFunction示例
     */
    public static void runCheckpointedFunctionExample(StreamExecutionEnvironment env) throws Exception {
        System.out.println("🚀 开始运行CheckpointedFunction示例...");
        
        // 创建计数器事件流
        DataStream<CounterEvent> counterEventStream = env.addSource(new MockCounterEventSource())
            .name("计数器事件源");
        
        // 1. 状态管理器示例
        DataStream<Tuple3<String, String, Long>> stateResultStream = counterEventStream
            .keyBy(event -> event.key)
            .map(new StateManager())
            .name("状态管理器");
        
        // 添加状态监控
        stateResultStream.addSink(new StateMonitorSink())
            .name("状态监控");
        
        stateResultStream.print("状态结果");
        
        // 2. 批处理状态管理器示例
        DataStream<Tuple2<String, Map<String, Long>>> batchResultStream = counterEventStream
            .keyBy(event -> event.key)
            .map(new BatchStateProcessor())
            .name("批处理状态管理器");
        
        // 过滤空结果
        batchResultStream.filter(result -> !result.f1.isEmpty())
            .print("批处理结果");
        
        System.out.println("✅ CheckpointedFunction示例配置完成");
        System.out.println("📸 提示: 状态会自动参与Checkpoint机制");
        System.out.println("🔄 提示: 作业重启时状态会自动恢复");
    }
    
    /**
     * 单独运行CheckpointedFunction示例的主方法
     */
    public static void main(String[] args) throws Exception {
        // 创建执行环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(2);
        
        // 配置状态后端
        com.flink.hbase.statemanagement.config.StateBackendConfig
            .configureDevRocksDB(env);
        
        // 运行示例
        runCheckpointedFunctionExample(env);
        
        // 执行作业
        env.execute("CheckpointedFunction示例");
    }
} 